#!/usr/bin/env python3
###
# @Author: Logan.Li
# @Gitee: https://gitee.com/attacker
# @email: admin@attacker.club
# @Date: 2025-03-14 05:17:55
# @LastEditTime: 2025-03-14 05:23:30
# @Description: MySQL backup and restore script
###

import os
import subprocess
from datetime import datetime
import pymysql
import argparse
import configparser
import sys

# デフォルト設定
OPERATION_MODE_DEFAULT = 'backup'
BACKUP_DIR_DEFAULT = "/root/backup/"
CONFIG_FILE = os.path.join(os.path.dirname(__file__), 'config.ini')

def load_config():
    if not os.path.exists(CONFIG_FILE):
        print(f"設定ファイルが見つかりません: {CONFIG_FILE}")
        sys.exit(1)
    
    config = configparser.ConfigParser()
    config.read(CONFIG_FILE)
    
    return {
        'backup': dict(config['backup']),
        'restore': dict(config['restore'])
    }

def get_databases(cursor):
    cursor.execute("SHOW DATABASES;")
    databases = cursor.fetchall()
    return [db[0] for db in databases if db[0] not in ['information_schema', 'performance_schema', 'mysql', 'sys']]

def do_backup(db_config, databases, backup_file_path):
    database_args = ' '.join([f'{db}' for db in databases])
    mysqldump_cmd = (
        f"mysqldump -u {db_config['user']} -p{db_config['password']} "
        f"-h {db_config['host']} -P {db_config['port']} "
        f"--single-transaction --quick --lock-tables=false "
        f"--set-gtid-purged=OFF --databases {database_args} > {backup_file_path}"
    )
    process = subprocess.run(mysqldump_cmd, shell=True, stderr=subprocess.PIPE)
    if process.returncode != 0:
        print(f"バックアップエラー: {process.stderr.decode()}")
        return False
    return True

def do_restore(db_config, restore_file):
    mysql_cmd = (
        f"mysql -u {db_config['user']} -p{db_config['password']} "
        f"-h {db_config['host']} -P {db_config['port']}"
    )
    
    process = subprocess.run(
        f"{mysql_cmd} < {restore_file}",
        shell=True,
        stderr=subprocess.PIPE
    )
    
    if process.returncode != 0:
        print(f"リストアエラー: {process.stderr.decode()}")
        return False
    return True

def main():
    parser = argparse.ArgumentParser(description='MySQL Database Backup and Restore Script')
    parser.add_argument('operation_mode', choices=['backup', 'restore'], default=OPERATION_MODE_DEFAULT, nargs='?',
                      help='Operation mode: "backup" or "restore"')
    parser.add_argument('--backup-file', dest='specified_backup_file',
                      help='Specify a backup file path for restore operation')
    parser.add_argument('--config', dest='config_file',
                      default=CONFIG_FILE,
                      help='Specify the configuration file path')
    
    args = parser.parse_args()
    
    try:
        configs = load_config()
        
        if args.operation_mode == 'backup':
            backup_file = f"{datetime.now().strftime('%Y%m%d%H%M%S')}_full_backup_exclude_sys.sql"
            backup_file_path = os.path.join(BACKUP_DIR_DEFAULT, backup_file)
            
            if not os.path.exists(BACKUP_DIR_DEFAULT):
                os.makedirs(BACKUP_DIR_DEFAULT)
            
            connection = pymysql.connect(**configs['backup'])
            cursor = connection.cursor()
            
            print("データベースリストを取得中...")
            databases_to_backup = get_databases(cursor)
            print(f"バックアップ対象: {', '.join(databases_to_backup)}")
            
            print("バックアップを開始します...")
            success = do_backup(configs['backup'], databases_to_backup, backup_file_path)
            
            if success:
                print(f"バックアップが完了しました。保存先: {backup_file_path}")
            else:
                print("バックアップが失敗しました。")
            
            cursor.close()
            connection.close()
        
        elif args.operation_mode == 'restore':
            specified_backup_file = args.specified_backup_file
            if not specified_backup_file:
                backup_files = [f for f in os.listdir(BACKUP_DIR_DEFAULT) if f.endswith('.sql')]
                if not backup_files:
                    print("バックアップファイルが見つかりません。")
                    sys.exit(1)
                
                specified_backup_file = os.path.join(
                    BACKUP_DIR_DEFAULT,
                    max(backup_files, key=lambda x: os.path.getctime(os.path.join(BACKUP_DIR_DEFAULT, x)))
                )
            
            if os.path.exists(specified_backup_file):
                print(f"リストアを開始します: {specified_backup_file}")
                success = do_restore(configs['restore'], specified_backup_file)
                
                if success:
                    print("リストアが完了しました。")
                else:
                    print("リストアが失敗しました。")
            else:
                print(f"指定されたバックアップファイルが存在しません: {specified_backup_file}")
    
    except Exception as e:
        print(f"エラーが発生しました: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()
            
            cursor.close()
            connection.close()
        
        elif args.operation_mode == 'restore':
            specified_backup_file = args.specified_backup_file
            if not specified_backup_file:
                backup_files = [f for f in os.listdir(BACKUP_DIR_DEFAULT) if f.endswith('.sql')]
                if not backup_files:
                    print("バックアップファイルが見つかりません。")
                    sys.exit(1)
                
                specified_backup_file = os.path.join(
                    BACKUP_DIR_DEFAULT,
                    max(backup_files, key=lambda x: os.path.getctime(os.path.join(BACKUP_DIR_DEFAULT, x)))
                )
            
            if os.path.exists(specified_backup_file):
                print(f"リストアを開始します: {specified_backup_file}")
                success = do_restore(configs['restore'], specified_backup_file)
                
                if success:
                    print("リストアが完了しました。")
                else:
                    print("リストアが失敗しました。")
            else:
                print(f"指定されたバックアップファイルが存在しません: {specified_backup_file}")
    
    except Exception as e:
        print(f"エラーが発生しました: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()